#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import numpy as np
import matplotlib.pyplot as plt
import torch
import glob
from tqdm import tqdm
import pandas as pd
import scipy.io
import re
import sys

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from data_provider.data_loader import ChannelDataset
from data_provider.data_factory import data_provider
from models import ChannelLLM
import argparse

def parse_args():
    parser = argparse.ArgumentParser(description='分析预测效果最差的样本')
    
    # 基本参数
    parser.add_argument('--results_dir', type=str, default='results', help='结果目录')
    parser.add_argument('--num_worst', type=int, default=20, help='要分析的最差样本数量')
    parser.add_argument('--model_path', type=str, default=None, help='模型路径，如果未指定则使用结果目录中的预测数据')
    parser.add_argument('--batch_size', type=int, default=4, help='批次大小')
    parser.add_argument('--n_clusters', type=int, default=25, help='簇数量')
    parser.add_argument('--n_probes', type=int, default=32, help='探头数量')
    
    return parser.parse_args()

def find_latest_result_dir(results_base_dir):
    """查找最新的结果目录"""
    result_dirs = glob.glob(os.path.join(results_base_dir, '*'))
    if not result_dirs:
        return None
    
    # 按修改时间排序
    result_dirs.sort(key=os.path.getmtime, reverse=True)
    return result_dirs[0]

def load_prediction_data(result_dir):
    """加载预测数据"""
    pred_path = os.path.join(result_dir, 'pred.npy')
    input_path = os.path.join(result_dir, 'input.npy')
    
    if not os.path.exists(pred_path) or not os.path.exists(input_path):
        print(f"找不到预测数据文件: {pred_path} 或 {input_path}")
        return None, None
    
    pred = np.load(pred_path)
    input_data = np.load(input_path)
    
    print(f"加载预测数据: {pred.shape}, 输入数据: {input_data.shape}")
    return pred, input_data

def load_dataset_time_ids():
    """加载数据集中的时间ID信息"""
    # 创建数据集实例来获取时间ID
    dataset = ChannelDataset(
        root_path='dataset/channel_data/',
        flag='test',
        size=[20, 10, 10],
        features='M',
        data_path='channel_data.npz',
        n_clusters=25,
        loader='channel',
        load_spatial_corr=True
    )
    
    # 返回时间ID信息
    return dataset.time_ids

def analyze_spatial_correlation_errors(pred, input_data, time_ids, num_worst=20):
    """
    分析空间相关性误差最大的样本
    
    Args:
        pred: 预测的探头权重, 形状 [B, seq_len, n_probes]
        input_data: 输入数据, 形状 [B, seq_len, n_clusters*5]
        time_ids: 时间点ID, 形状 [B, seq_len]
        num_worst: 要分析的最差样本数量
        
    Returns:
        worst_samples: 最差样本的统计信息
    """
    # 加载所有空间相关性数据以供比较
    spatial_corr_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'dy')
    
    # 检查目录是否存在
    if not os.path.exists(spatial_corr_dir):
        print(f"空间相关性数据目录不存在: {spatial_corr_dir}")
        return None
    
    # 加载所有空间相关性数据
    print(f"从 {spatial_corr_dir} 加载空间相关性数据...")
    spatial_corr_dict = {}
    spatial_files = glob.glob(os.path.join(spatial_corr_dir, 'spatial_correlation_*.mat'))
    
    for file_path in tqdm(spatial_files):
        match = re.search(r'spatial_correlation_(\d+)\.mat', os.path.basename(file_path))
        if match:
            time_id = int(match.group(1))
            try:
                mat_data = scipy.io.loadmat(file_path)
                if "p" in mat_data:
                    corr_data = mat_data['p'].flatten()
                    spatial_corr_dict[time_id] = np.concatenate([np.real(corr_data), np.imag(corr_data)])
            except Exception as e:
                print(f"加载文件 {file_path} 失败: {e}")
    
    print(f"加载了 {len(spatial_corr_dict)} 个空间相关性数据点")
    
    # 分析每个样本的误差
    batch_size, seq_len, n_probes = pred.shape
    
    # 保存每个样本的平均误差和最大误差
    sample_errors = []
    
    print("计算每个样本的误差...")
    for b in tqdm(range(batch_size)):
        for t in range(seq_len):
            # 获取当前时间点ID
            if time_ids is not None and b < len(time_ids) and t < len(time_ids[b]):
                time_id = time_ids[b, t]
            else:
                # 如果没有时间ID信息，使用索引作为标识
                time_id = f"batch_{b}_time_{t}"
            
            # 检查该时间点是否有空间相关性数据
            if time_id in spatial_corr_dict:
                true_spatial = spatial_corr_dict[time_id]
                
                # 获取预测的探头权重
                sample_pred = pred[b, t]
                
                # 由于我们没有重建的空间相关性，需要使用模型重建逻辑
                # 这里简化计算，使用与真实空间相关性的MSE作为误差
                if isinstance(true_spatial, np.ndarray) and true_spatial.size > 0:
                    # 如果你有重建函数，可以在这里应用
                    # reconstructed = reconstruct_spatial_corr(sample_pred)
                    # error = np.mean((reconstructed - true_spatial)**2)
                    
                    # 简化：记录时间点和探头权重直方图
                    error = 1.0  # 占位符
                    sample_errors.append({
                        'batch_idx': b,
                        'time_idx': t,
                        'time_id': time_id,
                        'error': error,
                        'pred_weights': sample_pred,
                        'true_spatial': true_spatial
                    })
    
    # 如果没有样本有误差记录，返回
    if not sample_errors:
        print("没有找到可以计算误差的样本")
        return None
    
    # 根据探头权重分布特征进行聚类分析
    print("对探头权重分布进行聚类分析...")
    
    # 提取所有样本的探头权重
    all_weights = np.array([s['pred_weights'] for s in sample_errors])
    
    # 计算每个样本的统计特征：
    # 1. 权重的熵 - 表示分布的平坦程度
    # 2. 最大权重值 - 表示是否有明显主导探头
    # 3. 非零权重比例 - 表示激活的探头数量
    
    features = []
    for weights in all_weights:
        # 归一化权重确保和为1
        norm_weights = weights / (np.sum(weights) + 1e-10)
        
        # 计算熵
        entropy = -np.sum(norm_weights * np.log2(norm_weights + 1e-10))
        
        # 最大权重
        max_weight = np.max(norm_weights)
        
        # 有效探头比例 (权重>0.01的探头)
        active_probes = np.sum(norm_weights > 0.01) / len(norm_weights)
        
        features.append([entropy, max_weight, active_probes])
    
    features = np.array(features)
    
    # 使用K均值聚类
    from sklearn.cluster import KMeans
    n_clusters = min(5, len(features))
    kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(features)
    clusters = kmeans.labels_
    
    # 将聚类结果添加到样本信息中
    for i, s in enumerate(sample_errors):
        s['cluster'] = int(clusters[i])
    
    # 按聚类统计结果
    cluster_stats = {}
    for cluster_id in range(n_clusters):
        cluster_samples = [s for s in sample_errors if s['cluster'] == cluster_id]
        cluster_stats[cluster_id] = {
            'count': len(cluster_samples),
            'time_ids': [s['time_id'] for s in cluster_samples[:10]]  # 只显示前10个
        }
    
    print("\n聚类统计:")
    for cluster_id, stats in cluster_stats.items():
        print(f"  聚类 {cluster_id}: {stats['count']} 个样本")
        print(f"    代表性时间ID: {stats['time_ids'][:5]}...")
    
    # 将样本按时间ID分组
    time_id_groups = {}
    for s in sample_errors:
        time_id = s['time_id']
        if time_id not in time_id_groups:
            time_id_groups[time_id] = []
        time_id_groups[time_id].append(s)
    
    # 分析每个时间点的情况
    time_stats = []
    for time_id, samples in time_id_groups.items():
        # 计算该时间点的平均特征
        avg_features = np.mean([features[i] for i in range(len(sample_errors)) if sample_errors[i]['time_id'] == time_id], axis=0)
        time_stats.append({
            'time_id': time_id,
            'count': len(samples),
            'avg_entropy': avg_features[0],
            'avg_max_weight': avg_features[1],
            'avg_active_probes': avg_features[2]
        })
    
    # 按出现频率排序
    time_stats.sort(key=lambda x: x['count'], reverse=True)
    
    print("\n时间点出现频率统计 (前10个):")
    for ts in time_stats[:10]:
        print(f"  时间点 {ts['time_id']}: 出现 {ts['count']} 次, "
              f"平均熵: {ts['avg_entropy']:.4f}, "
              f"平均最大权重: {ts['avg_max_weight']:.4f}, "
              f"平均活跃探头比例: {ts['avg_active_probes']:.4f}")
    
    # 创建可视化
    print("\n生成图表...")
    
    # 1. 聚类散点图
    plt.figure(figsize=(10, 8))
    plt.scatter(features[:, 0], features[:, 1], c=clusters, cmap='viridis', alpha=0.6)
    plt.xlabel('熵 (分布平坦程度)')
    plt.ylabel('最大权重')
    plt.title('探头权重分布特征聚类')
    plt.colorbar(label='聚类ID')
    plt.savefig('probe_weights_clusters.png')
    
    # 2. 最常出现的时间点的权重分布
    if len(time_stats) > 0:
        top_time_id = time_stats[0]['time_id']
        top_samples = time_id_groups[top_time_id]
        
        plt.figure(figsize=(12, 6))
        for i, s in enumerate(top_samples[:5]):  # 只显示前5个样本
            plt.subplot(1, 5, i+1)
            plt.bar(range(len(s['pred_weights'])), s['pred_weights'])
            plt.title(f"样本 {i+1}")
            plt.xlabel('探头索引')
            plt.ylabel('权重')
        plt.tight_layout()
        plt.savefig(f'time_id_{top_time_id}_weights.png')
    
    # 返回分析结果
    return {
        'time_stats': time_stats,
        'cluster_stats': cluster_stats
    }

def main():
    args = parse_args()
    
    # 查找最新的结果目录
    if args.results_dir == 'results':
        args.results_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'results')
    
    result_dir = find_latest_result_dir(args.results_dir)
    if result_dir is None:
        print(f"找不到结果目录: {args.results_dir}")
        return
    
    print(f"使用结果目录: {result_dir}")
    
    # 加载预测数据
    pred, input_data = load_prediction_data(result_dir)
    if pred is None or input_data is None:
        return
    
    # 加载时间ID信息
    print("加载时间ID信息...")
    time_ids = load_dataset_time_ids()
    
    # 分析误差最大的样本
    print("分析预测效果最差的样本...")
    analysis_results = analyze_spatial_correlation_errors(
        pred, input_data, time_ids, num_worst=args.num_worst
    )
    
    if analysis_results:
        print("\n分析完成！")
        print(f"生成的图表已保存在当前目录")
        
        # 保存分析结果到CSV
        if analysis_results['time_stats']:
            df = pd.DataFrame(analysis_results['time_stats'])
            csv_path = 'time_id_analysis.csv'
            df.to_csv(csv_path, index=False)
            print(f"时间点分析结果已保存到: {csv_path}")

if __name__ == "__main__":
    main() 